package main

import (
	"SQLi/model"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"xorm.io/xorm"
)

type Database struct {
	engine *xorm.Engine
}

func NewDatabase(username, password string) (*Database, error) {
	dsn := fmt.Sprintf("%s:%s@tcp(127.0.0.1:3306)/go_sec_labs?charset=utf8mb4&parseTime=True&loc=Local", username, password)
	engine, err := xorm.NewEngine("mysql", dsn)
	if err!= nil {
		return nil, fmt.Errorf("failed to connect to database: %v", err)
	}
	engine.ShowSQL(true) // 显示生成的 SQL 语句，便于调试
	return &Database{engine: engine}, nil
}

func (db *Database) Close() error {
	return db.engine.Close()
}

func (db *Database) SelectUsersByName(name string) []model.User {
	var users []model.User
	//err := db.engine.Where(fmt.Sprintf("name = '%s'", name)).Find(&users)
	err := db.engine.Where("name = ?", name).Find(&users)
	if err!= nil {
		log.Printf("Error %v\n", err)
		return nil
	}
	return users
}

func (db *Database) SelectUsersLimit(limit, offset int) []model.User {
	var users []model.User
	err := db.engine.Limit(limit, offset).Find(&users)
	if err!= nil {
		log.Printf("Error %v\n", err)
		return nil
	}
	return users
}

func (db *Database) SelectUsersIn(names []string) []model.User {
	var users []model.User
	err := db.engine.In("name", names).Find(&users)
	//str := "'" + strings.Join(names, "','") + "'"
	//err := db.engine.Where(fmt.Sprintf("name in (%s)", str)).Find(&users)
	if err!= nil {
		log.Printf("Error %v\n", err)
		return nil
	}
	return users
}

func (db *Database) SelectUsersLike(name string) []model.User {
	var users []model.User
	//err := db.engine.Where("name like '%" + name + "%'").Find(&users)
	err := db.engine.Where("name like ?", "%"+name+"%").Find(&users)
	if err!= nil {
		log.Printf("Error %v\n", err)
		return nil
	}
	return users
}

func (db *Database) SelectUsersOrderBy(field string) []model.User {
	// 定义允许的排序字段白名单
	allowedFields := map[string]bool{
		"id":    true,
		"name":  true,
		"age":   true,
		"sex":   true,
	}
	// 检查用户提供的排序字段是否在白名单中
	if !allowedFields[field] {
		log.Println("Invalid order field:", field)
		return nil
	}

	var users []model.User
	err := db.engine.Desc(field).Find(&users)
	if err!= nil {
		log.Printf("Error %v\n", err)
		return nil
	}
	return users
}

func (db *Database) InsertBatch(users []model.User) int64 {
	affected, err := db.engine.Insert(users)
	if err!= nil {
		log.Printf("Error %v\n", err)
		return 0
	}
	return affected
}

func main() {
	// 初始化数据库
	db, err := NewDatabase("root", "111111")
	if err!= nil {
		log.Fatalf("Database initialization failed: %v", err)
	}
	defer db.Close()

	// 查询用户
	//users := db.SelectUsersByName("' or 1=1) -- ")
	//users := db.SelectUsersIn([]string{"zhangsan", "wangwu') or 1=1) -- "})
	//users := db.SelectUsersLike("an%' or 1=1) -- ")
	//users := db.SelectUsersOrderBy("id`,if(substr(database(),1,1)='a',sleep(2),age),`id")
	//users := db.SelectUsersOrderBy("age")
	affected := db.InsertBatch([]model.User{
		{Name: "'+if(substr(database(),1,1)='g',sleep(3),1)+'", Age: 44, Sex: "0", Password: "12121"},
		{Name: "bb", Age: 66, Sex: "1", Password: "32323"},
	})
	//for _, user := range users {
	//	fmt.Println(user)
	//}
	fmt.Println(affected)
}